OUT_DIR := build
INC_DIR := $(realpath ../../include)
LIB_DIR := $(realpath ../../build/lib)
LLVM_PRJ_DIR := $(realpath ../../3rdparty/llvm)

LLVM_CLANG_INC_DIR := $(LLVM_PRJ_DIR)/build/lib/clang/17/include
DESIGN_JSON_FILE := $(realpath ./design2.json)

CC := /usr/bin/gcc
CXX := /usr/bin/g++

WAFER_OUT_DIR := $(realpath ../../build)
WAFER_FRONTEND := $(WAFER_OUT_DIR)/bin/wafer-frontend
WAFER_OPT := $(WAFER_OUT_DIR)/bin/wafer-opt
WAFER_TRANSLATE := $(WAFER_OUT_DIR)/bin/wafer-translate

CLANGXX := $(LLVM_PRJ_DIR)/build/bin/clang++
LLC := $(LLVM_PRJ_DIR)/build/bin/llc

LLVM_LIB_DIR := $(LLVM_PRJ_DIR)/build/lib

CUDA_BASE_DIR := /usr/local/cuda
CUDA_INC_DIR := $(CUDA_BASE_DIR)/targets/x86_64-linux/include
CUDA_LIB_DIR := $(CUDA_BASE_DIR)/targets/x86_64-linux/lib

CUDAFEXX := $(CUDA_BASE_DIR)/bin/cudafe++
CICC := $(CUDA_BASE_DIR)/nvvm/bin/cicc
PTXAS := $(CUDA_BASE_DIR)/bin/ptxas
FATBINARY := $(CUDA_BASE_DIR)/bin/fatbinary
NVLINK := $(CUDA_BASE_DIR)/bin/nvlink
CUDA_CRT_STUB := $(CUDA_BASE_DIR)/bin/crt/link.stub

OPENSSL_OUT_DIR := $(realpath ../../3rdparty/openssl/build)
OPENSSL_LIB_DIR := $(OPENSSL_DIR)/lib64
OPENSSL_INC_DIR := $(OPENSSL_DIR)/include


export PRINT_SCHEDULE=ON


.PHONY: all
all: $(OUT_DIR)/host_ctrl.o $(OUT_DIR)/AESExample.o $(OUT_DIR)/wafer_gpu.out
	$(CLANGXX) $(OUT_DIR)/host_ctrl.o $(OUT_DIR)/AESExample.o \
		-L $(LLVM_LIB_DIR) -lmlir_c_runner_utils \
		-L $(LIB_DIR) -lwafer_sst_runtime \
		-L $(LIB_DIR)/CGRAAPI -lwafer_cgra_api \
		-Wl,-rpath,$(LLVM_LIB_DIR):$(LIB_DIR) \
		-o $(OUT_DIR)/aes_crack
	@echo "Build executable (w/GPU) successfully"

.PHONY: cpu-only
cpu-only: $(OUT_DIR)/host_ctrl.o $(OUT_DIR)/AESExample.o
	$(CLANGXX) $(OUT_DIR)/host_ctrl.o $(OUT_DIR)/AESExample.o \
		-L $(LLVM_LIB_DIR) -lmlir_c_runner_utils -lmlir_runner_utils \
		-L $(LIB_DIR) -lwafer_sst_runtime \
		-L $(LIB_DIR)/CGRAAPI -lwafer_cgra_api \
		-Wl,-rpath,$(LLVM_LIB_DIR):$(LIB_DIR) \
		-o $(OUT_DIR)/aes_crack
	@echo "Build (CPU Only) successfully"

.PHONY: run
run: all
	./build/aes_crack

$(OUT_DIR)/host_ctrl.o: $(OUT_DIR)/op_aes.ll
	$(LLC) -mtriple=x86_64-unknown-linux-gnu $< --filetype=obj -o $@


$(OUT_DIR)/op_aes.ll: $(OUT_DIR)/aescrack_op.mlir
	DESIGN_JSON_FILE=${DESIGN_JSON_FILE} $(WAFER_OPT) $^ \
		--llvm-request-c-wrappers \
		--convert-scf-to-cf \
		--convert-arith-to-llvm \
		--finalize-memref-to-llvm \
		--convert-vector-to-llvm \
		--lower-sst-to-llvm \
		--convert-func-to-llvm \
		--convert-index-to-llvm \
		--reconcile-unrealized-casts | \
	$(WAFER_TRANSLATE) --wafer-to-llvmir -o $@


$(OUT_DIR)/aescrack_op.mlir: $(OUT_DIR)/aescrack_sst.mlir
	DESIGN_JSON_FILE=${DESIGN_JSON_FILE} $(WAFER_OPT) $^ \
		--sst-lower-hyper \
		--lower-crypto-padmessage \
		--lower-cryptosha1-without-SHA1ISA \
		--lower-crypto-to-cgra -o $@



$(OUT_DIR)/aescrack_sst.mlir: $(OUT_DIR)/aescrack_scheduled.mlir
	DESIGN_JSON_FILE=${DESIGN_JSON_FILE} $(WAFER_OPT) $^ \
		--lower-hyper-sstmemcpy -o $@

$(OUT_DIR)/aescrack_scheduled.mlir: $(OUT_DIR)/aescrack_hyper.mlir
	DESIGN_JSON_FILE=${DESIGN_JSON_FILE} $(WAFER_OPT) $^ \
		--sst-device-schedule -o $@



$(OUT_DIR)/aescrack_hyper.mlir: $(OUT_DIR)/aescrack.mlir
	DESIGN_JSON_FILE=${DESIGN_JSON_FILE} $(WAFER_OPT) $^ --lower-crypto-to-hyper -o $@

$(OUT_DIR)/aescrack.mlir: AESExample.cpp
	@mkdir -p $(@D)
	$(WAFER_FRONTEND) -I $(INC_DIR) -I $(LLVM_CLANG_INC_DIR) \
		 -c $< -o $@

# another way to get op_aes.ll:
# $(OUT_DIR)/op_aes.ll: AESExample.cpp
# 	mkdir -p $(@D)
# 	$(WAFER_FRONTEND) -I $(INC_DIR) -I $(LLVM_CLANG_INC_DIR) \
# 		--sst --JsonFilePath $(DESIGN_JSON_FILE) -c $< --wafer-to-llvmir -o $@

$(OUT_DIR)/%.o: %.cpp
	@mkdir -p $(@D)
	$(CLANGXX) -c $< -I$(INC_DIR) -o $@

# cuda 12.3
# make CUDA_BASE_DIR=/usr/local/cuda-12.3
# Generated with command `nvcc main.cu -dryrun`, patched only with directory and filename rewrite.
# nvcc using cuda 12.3 version
$(OUT_DIR)/wafer_gpu.out: $(OUT_DIR)/kernel.ptx main.cu
	$(CC) -D__CUDA_ARCH__=520 -D__CUDA_ARCH_LIST__=520 -E -x c++ \
		-DCUDA_DOUBLE_MATH_FUNCTIONS -D__CUDACC__ -D__NVCC__ \
		"-I$(CUDA_INC_DIR)" -D__CUDACC_VER_MAJOR__=12 -D__CUDACC_VER_MINOR__=3 -D__CUDACC_VER_BUILD__=107 -include "cuda_runtime.h" \
		-m64 "./main.cu" > "$(OUT_DIR)/main.cpp1.ii"
	$(CICC) --c++17 --gnu_version=110400 --display_error_number --allow_managed \
		-arch compute_52 -m64 -ftz=0 -prec_div=1 -prec_sqrt=1 -fmad=1 --include_file_name "main.fatbin.c" -tused -nvvmir-library \
		"$(CUDA_BASE_DIR)/nvvm/libdevice/libdevice.10.bc" --gen_module_id_file --module_id_file_name "$(OUT_DIR)/main.module_id" --orig_src_file_name "./main.cu" \
		--gen_c_file_name "$(OUT_DIR)/main.cudafe1.c" --stub_file_name "$(OUT_DIR)/main.cudafe1.stub.c" \
		--gen_device_file_name "$(OUT_DIR)/main.cudafe1.gpu" "$(OUT_DIR)/main.cpp1.ii" -o "$(OUT_DIR)/main.ptx"
	cp $(OUT_DIR)/kernel.ptx $(OUT_DIR)/main.ptx
	$(FATBINARY) --create="$(OUT_DIR)/main.fatbin" -64 "--image=profile=compute_52,file=$(OUT_DIR)/main.ptx" --embedded-fatbin="$(OUT_DIR)/main.fatbin.c"
	$(CC) -E -x c++ -D__CUDACC__ -D__NVCC__  "-I$(CUDA_INC_DIR)" -D__CUDACC_VER_MAJOR__=12 -D__CUDACC_VER_MINOR__=3 -D__CUDACC_VER_BUILD__=107 \
		-include "cuda_runtime.h" -m64 "./main.cu" > "$(OUT_DIR)/main.cpp4.ii"
	$(CUDAFEXX) --c++17 --gnu_version=110400 --display_error_number --allow_managed  --m64 --parse_templates \
		--gen_c_file_name "$(OUT_DIR)/main.cudafe1.cpp" --stub_file_name "main.cudafe1.stub.c" \
		--module_id_file_name "$(OUT_DIR)/main.module_id" "$(OUT_DIR)/main.cpp4.ii"
	$(CC) -D__CUDA_ARCH__=520 -c -x c++ -DCUDA_DOUBLE_MATH_FUNCTIONS "-I$(CUDA_INC_DIR)" -m64 -o "$(OUT_DIR)/main.o" "$(OUT_DIR)/main.cudafe1.cpp"
	$(NVLINK) --arch=compute_52 --register-link-binaries="$(OUT_DIR)/dlink.reg.c"  -m64 "-L$(CUDA_LIB_DIR)/stubs" "-L$(CUDA_LIB_DIR)" -cpu-arch=X86_64 "$(OUT_DIR)/main.o"
	$(FATBINARY) --create="$(OUT_DIR)/dlink.fatbin" -64 -link --embedded-fatbin="$(OUT_DIR)/dlink.fatbin.c"
	$(CC) -c -x c++ -DFATBINFILE="\"dlink.fatbin.c\"" -DREGISTERLINKBINARYFILE="\"dlink.reg.c\"" -I$(OUT_DIR) \
		-D__NV_EXTRA_INITIALIZATION= -D__NV_EXTRA_FINALIZATION= -D__CUDA_INCLUDE_COMPILER_INTERNAL_HEADERS__ \
		"-I$(CUDA_INC_DIR)" -D__CUDACC_VER_MAJOR__=12 -D__CUDACC_VER_MINOR__=3 -D__CUDACC_VER_BUILD__=107 \
		-m64 -o "$(OUT_DIR)/dlink.o" "$(CUDA_BASE_DIR)/bin/crt/link.stub"
	$(CXX) -m64 -o $@ -Wl,--start-group "$(OUT_DIR)/dlink.o" "$(OUT_DIR)/main.o" "-L$(CUDA_LIB_DIR)/stubs" "-L$(CUDA_LIB_DIR)" -lcudadevrt  -L$(CUDA_BASE_DIR)/lib64 -lcudart_static -lrt -lpthread -ldl  -Wl,--end-group

$(OUT_DIR)/kernel.ptx: $(OUT_DIR)/gpulaunch.mlir
	@mkdir -p $(@D)
	$(WAFER_OPT) $< -device-schedule -lower-hyper-memcpy -lower-hyper | \
		$(WAFER_OPT) -lower-crypto-padmessage | \
		$(WAFER_OPT) -lower-cryptosha1-without-SHA1ISA | \
		$(WAFER_OPT) -convert-async-to-llvm | \
		$(WAFER_OPT) -llvm-request-c-wrappers | \
		$(WAFER_OPT) -convert-scf-to-cf | \
		$(WAFER_OPT) -gpu-kernel-outlining | \
		$(WAFER_OPT) -convert-arith-to-llvm | \
		$(WAFER_OPT) -finalize-memref-to-llvm | \
		$(WAFER_OPT) -convert-vector-to-llvm | \
		$(WAFER_OPT) -pass-pipeline='builtin.module(gpu.module(strip-debuginfo,convert-gpu-to-nvvm,reconcile-unrealized-casts,gpu-to-cubin{chip=sm_30 features=+ptx32}))' > /dev/null
	@mv ./kernel.ptx $@

$(OUT_DIR)/gpulaunch.mlir: AESExample.cpp $(DESIGN_JSON_FILE)
	@mkdir -p $(@D)
	$(WAFER_FRONTEND) -I $(INC_DIR) -I $(LLVM_CLANG_INC_DIR) \
		--ptx --JsonFilePath $(DESIGN_JSON_FILE) -c $< -o $@
		
.PHONY: openssl
openssl: AESExample.cpp
	@mkdir -p build
	$(CXX) -DWAFER_OPENSSL $< -I $(OPENSSL_INC_DIR) -L $(OPENSSL_LIB_DIR) -l crypto -o build/aes_crack

.PHONY: clean
clean:
	rm -rf $(OUT_DIR)
